# -*-coding:utf-8-*-

# 获取序列window
import pandas as pd
import re
import random

windows = 31


def find_all(sub, s):
    index_list = []
    index = s.find(sub)
    while index != -1:
        index_list.append(index)
        index = s.find(sub, index + 1)

    if len(index_list) > 0:
        return index_list
    else:
        return -1


def findSequence(file: str, findsequences: str, windows: int = 15, site: int = 0) -> str:
    '''
    获得序列框内序列
    :param file: string  文件名
    :param findsequences: string 需要比对序列
    :param windows:  int  序列框大小，算一半，默认为15
    :param site: int 位置，默认为0
    :return: string 序列
    '''

    file = open(file, 'r')
    stringfile = file.readlines()
    singlestringfile = ''
    count = True
    for j in stringfile:
        if count:
            count = False
        else:
            singlestringfile = singlestringfile + j
    # 合成一列
    singlestringfile = singlestringfile.replace('\n', '')
    # ISNVNK(1)ALDFIASK
    acetyle = findsequences
    # 搜寻乙酰化位点
    # print(re.findall('\([\d.]+\)', acetyle))
    # print(re.search('\([\d.]+\)', acetyle).span())
    # print(find_all('(1)', acetyle))
    acetyles = find_all('(1)', acetyle)
    if acetyles == -1:
        acetyleFirst, acetyleLast = re.search('\([\d.]+\)', acetyle).span()
    elif len(acetyles) > 1:
        acetyleFirst = acetyles[site]
        acetyleLast = acetyles[site] + 3
    else:
        acetyleFirst, acetyleLast = re.search('\([\d.]+\)', acetyle).span()
    # 拼接除去(1)的序列
    finallsite = re.sub('\([\d.]+\)', '', acetyle)
    # finallsite = acetyle[:acetyleFirst] + acetyle[acetyleLast:]
    # 找到在总序列里面的位置
    sequenceFirst, sequenceLast = re.search(finallsite, singlestringfile).span()
    # 阅读框的起位置
    windowsFirst = sequenceFirst + acetyleFirst - 1 - windows
    # 阅读框的结尾位置
    windowsLast = sequenceFirst + acetyleFirst + windows
    # 拼接为数据
    if windowsFirst < 0:
        windowsdata = singlestringfile[0:windowsLast]
    else:
        windowsdata = singlestringfile[windowsFirst:windowsLast]
    # print(windowsdata)
    # 两端补齐
    if windowsFirst < 0:
        windowsdata = '-' * abs(windowsFirst) + windowsdata
    if len(singlestringfile) - windowsLast < 0:
        windowsdata = windowsdata + '-' * abs(len(singlestringfile) - windowsLast)
    return windowsdata

# 阅读初始文件，为acdata.xlsx文件夹
data = pd.read_excel('acdata.xlsx')
#去除里面的重复选项
fastalist = []
#序列组
sequences = []
labels = []

sql = open('makesumsequences.fasta', 'w')
total = open('total.csv', 'w')
totalnum = 0
for i in range(len(data)):
    ratio = data['combined Ratio M/L'][i]
    if ratio == 'NA':
        continue
    elif ratio > 1.2:
        labels.append(1)
    elif ratio < 0.9 or ratio > 1.1:
        continue
    else:
        labels.append(0)

    headers = data['headers'][i]
    headers = headers.split('|')
    # print(headers[1])
    site = 1
    if headers[1] not in fastalist:
        fastalist.append(headers[1])
        site = 0
    try:
        findsq = findSequence(f'fastadata/{headers[1]}.fasta', data['Acetyl..K..Probabilities'][i], 15, site)
        sequences.append(findsq)

        # 保存总文件
        pdbname = f'>tr-' + str(totalnum).rjust(5, '0')
        sql.write(pdbname)
        sql.write('\n')
        sql.write(re.sub('[^ARNDCQEGHILKMFPSTWYV-]', '-', ''.join(findsq).upper()))
        sql.write('\n')

        # 保存单个文件
        f2 = open(f'fasta2data/{pdbname[4:]}.fasta', 'w')
        f2.write(pdbname)
        f2.write('\n')
        f2.write(re.sub('[^ARNDCQEGHILKMFPSTWYV-]', '-', ''.join(findsq).upper()))
        f2.close()

        # 编号名字
        total.write(pdbname[4:])
        total.write(',')
        total.write(pdbname)
        total.write(',')
        # 标签名字
        total.write(headers[1])
        # 数值
        total.write(',')
        total.write(str(data['combined Ratio M/L'][i]))
        # 标签
        total.write(',')
        if ratio > 1.2:
            total.write('1')
        else:
            total.write('0')
        total.write(',')
        # 序列
        total.write(re.sub('[^ARNDCQEGHILKMFPSTWYV-]', '-', ''.join(findsq).upper()))
        total.write(',')
        total.write(str(data['combined Ratios H/L'][i]))
        total.write('\n')

        totalnum = totalnum + 1

    except Exception as e:
        print(headers[1])
        labels.pop()
        print(str(e))
sql.close()
total.close()
print(len(sequences))
print(len(labels))
# a = findSequence('fastadata/Q67FY2.fasta', 'NPQAGVSPFSSLK(1)GK(1)VK', 15)
# print(a)
# print(sequences)

#
# import sys
#
# sys.path.append('..')
# from CKSAAP import CKSAAP
# import numpy as np
#
# cksaap = []
# newsequences = []
# for i in sequences:
#     sequence = re.sub('[^ARNDCQEGHILKMFPSTWYV-]', '-', ''.join(i).upper())
#     cksaap.append(CKSAAP(sequence, 1))
#     newsequences.append(sequence)
#
# X = np.array(cksaap)
# y = np.array(labels)
# Z = np.array(newsequences)
# print(Z.shape)

# from sklearn.linear_model import LogisticRegression
#
# print(X)
# print(y)
#


########################################################################
# 保存sequences
# f = open('sequences.txt', 'w')
# for i in newsequences:
#     f.write(i)
#     f.write(',')
# f.close()
# f2 = open('labels.txt', 'w')
# for j in labels:
#     f2.write(str(j))
#     f2.write(',')
# f2.close()
#########################################################################

# np.savetxt('sequences.txt', Z, delimiter=',')
# np.savetxt('labels.txt', y, delimiter=',')
#
# from sklearn.model_selection import train_test_split
#
# X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2,
#                                                     stratify=y,
#                                                     shuffle=True,
#                                                     random_state=1)
#
# clf = LogisticRegression().fit(X_train, y_train)
# score = clf.score(X_test, y_test)
# print(score)
#
# from sklearn.ensemble import RandomForestClassifier
#
# clf = RandomForestClassifier(max_depth=2).fit(X_train, y_train)
# score = clf.score(X_test, y_test)
# print(score)
